import os
import numpy as np
import SimpleITK as sitk
from tqdm import tqdm
import logging

# === 配置日志 ===
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# === 预处理函数 ===
def preprocess_volume(slice_data):
    """预处理单个CT切片"""
    # 阈值截取
    slice_data = np.clip(slice_data, -200, 200)
    slice_data = slice_data.astype(np.float32)
    slice_data = slice_data / 200
    return slice_data

# === 肝脏切片定位函数 ===
def find_liver_slices(mask_array):
    """找到肝脏区域开始和结束的slice，并各向外扩张slice"""
    z = np.any(mask_array, axis=(1, 2))
    start_slice, end_slice = np.where(z)[0][[0, -1]]
    start_slice = max(0, start_slice - 1)
    end_slice = min(mask_array.shape[0] - 1, end_slice + 2)
    return start_slice, end_slice

# === 中心裁剪函数 ===
def crop_center(img, croph, cropw):
    """从图像中心裁剪出指定大小的区域"""
    height, width = img[0].shape
    starth = height//2 - (croph//2)
    startw = width//2 - (cropw//2)
    return img[:, starth:starth+croph, startw:startw+cropw]

# === 单个体积处理函数 ===
def process_volume(ct_path, mask_path, output_img_dir, output_mask_dir, index):
    """处理单个CT体积及其对应的mask"""
    try:
        # 加载CT和mask文件
        ct_src = sitk.ReadImage(ct_path, sitk.sitkInt16)
        mask = sitk.ReadImage(mask_path, sitk.sitkUInt8)
        
        # 转换为数组
        ct_array = sitk.GetArrayFromImage(ct_src)
        mask_array = sitk.GetArrayFromImage(mask)

        # 预处理整个体积
        ct_array = preprocess_volume(ct_array)

        # 定位肝脏切片范围
        start_slice, end_slice = find_liver_slices(mask_array)
        logger.info(f"处理 {os.path.basename(ct_path)}: 切片范围 {start_slice}-{end_slice}")

        # 裁剪肝脏区域
        ct_crop = ct_array[start_slice:end_slice, :, :]
        mask_crop = mask_array[start_slice+1:end_slice-1, :, :]

        # 空间裁剪 (448x448)
        ct_crop = crop_center(ct_crop, 448, 448)
        mask_crop = crop_center(mask_crop, 448, 448)

        # 切片处理
        valid_slices = 0
        for n_slice in range(mask_crop.shape[0]):
            if np.sum(mask_crop[n_slice]) > 0:  # 仅保存有病灶的切片
                # 构建三通道输入
                ct_slice_data = np.stack([
                    ct_crop[n_slice],
                    ct_crop[n_slice + 1],
                    ct_crop[n_slice + 2]
                ], axis=-1).astype(np.float32)

                # 保存切片
                img_path = os.path.join(output_img_dir, f"{index+1}_{n_slice}.npy")
                mask_path = os.path.join(output_mask_dir, f"{index+1}_{n_slice}.npy")
                
                np.save(img_path, ct_slice_data)
                np.save(mask_path, mask_crop[n_slice].astype(np.uint8))
                valid_slices += 1
        
        return valid_slices
    except Exception as e:
        logger.error(f"处理 {ct_path} 失败: {str(e)}")
        return 0

# === 主处理函数 ===
def process_dataset(ct_dir, mask_dir, output_img_dir, output_mask_dir, dataset_name="数据集"):
    """处理整个数据集"""
    os.makedirs(output_img_dir, exist_ok=True)
    os.makedirs(output_mask_dir, exist_ok=True)
    
    ct_files = [f for f in os.listdir(ct_dir) if f.endswith('.nii')]
    total_slices = 0
    
    with tqdm(ct_files, desc=f"处理{dataset_name}", unit="文件") as pbar:
        for idx, ct_file in enumerate(pbar):
            # 构建mask文件名
            mask_file = ct_file.replace("volume", "segmentation")
            ct_path = os.path.join(ct_dir, ct_file)
            mask_path = os.path.join(mask_dir, mask_file)
            
            if not os.path.exists(mask_path):
                logger.warning(f"Mask文件不存在: {mask_path}, 跳过")
                continue
                
            # 处理单个体积
            slices_saved = process_volume(
                ct_path, mask_path, 
                output_img_dir, output_mask_dir, 
                idx
            )
            total_slices += slices_saved
            pbar.set_postfix_str(f"已保存切片: {slices_saved}")
    
    logger.info(f"{dataset_name}处理完成! 共保存 {total_slices} 个有效切片")

# === 主程序 ===
if __name__ == "__main__":
    # 配置路径
    input_path = 'LITS_dataset'
    output_path = ''
    
    # 训练集路径
    ct_train = os.path.join(input_path, 'train/data')
    seg_train = os.path.join(input_path, 'train/label')
    img_train_out = os.path.join(output_path, 'trainImage')
    mask_train_out = os.path.join(output_path, 'trainMask')
    
    # 测试集路径
    ct_test = os.path.join(input_path, 'test/data')
    seg_test = os.path.join(input_path, 'test/label')
    img_test_out = os.path.join(output_path, 'testImage')
    mask_test_out = os.path.join(output_path, 'testMask')
    
    # 处理训练集
    process_dataset(ct_train, seg_train, img_train_out, mask_train_out, "训练集")
    
    # 处理测试集
    process_dataset(ct_test, seg_test, img_test_out, mask_test_out, "测试集")
    
    logger.info("所有处理完成!")
